import sqlite3

class Database:
    def __init__(self):
        self.conn = None
        self.cursor = None
    
    def connect(self, db_name='rabbit_info.db'):
        """连接数据库"""
        self.conn = sqlite3.connect(db_name)
        self.cursor = self.conn.cursor()
    
    def init_database(self):
        """初始化数据库表结构"""
        if not self.cursor:
            self.connect()
        
        # 创建信息表（移除type字段，保留category字段并将其用于标签）
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS information (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                user_id INTEGER NOT NULL,
                title TEXT NOT NULL,
                content TEXT,
                category TEXT,  -- 改为标签字段
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )
        ''')
        
        # 创建用户表
        self.cursor.execute('''
            CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT NOT NULL UNIQUE,
                password_hash TEXT NOT NULL,
                email TEXT NOT NULL UNIQUE,
                is_verified BOOLEAN DEFAULT FALSE,
                verification_code TEXT,
                reset_token TEXT,
                reset_token_expiry TIMESTAMP,
                reset_verification_code TEXT,
                create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                last_login TIMESTAMP
            )
        ''')
        
        # 为每个用户自动创建默认的"未分类"类型
        # 这里我们不在初始化时插入，而是在用户首次使用时动态创建
        
        # 更新现有数据库表结构，确保所有必要的列都存在
        self.update_database_structure()
        
    def update_database_structure(self):
        """更新数据库表结构，添加缺失的列，移除或修改不需要的字段"""
        try:
            # 检查并添加reset_verification_code列（如果不存在）
            self.cursor.execute("PRAGMA table_info(users)")
            columns = [column[1] for column in self.cursor.fetchall()]
            
            if 'reset_verification_code' not in columns:
                self.cursor.execute("ALTER TABLE users ADD COLUMN reset_verification_code TEXT")
                self.conn.commit()
            
            # 检查并添加user_id列到information表（如果不存在）
            self.cursor.execute("PRAGMA table_info(information)")
            info_columns = [column[1] for column in self.cursor.fetchall()]
            
            if 'user_id' not in info_columns:
                self.cursor.execute("ALTER TABLE information ADD COLUMN user_id INTEGER DEFAULT 0")
                self.conn.commit()
            
            # 检查并处理type字段
            if 'type' in info_columns:
                # 尝试修改type字段，移除NOT NULL约束
                try:
                    # SQLite不支持直接修改列约束，所以我们需要重命名表并重新创建
                    self.cursor.execute("ALTER TABLE information RENAME TO information_old")
                    
                    # 创建新表结构
                    self.cursor.execute('''
                        CREATE TABLE information (
                            id INTEGER PRIMARY KEY AUTOINCREMENT,
                            user_id INTEGER NOT NULL,
                            title TEXT NOT NULL,
                            content TEXT,
                            category TEXT,  -- 标签字段
                            create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                            update_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                            FOREIGN KEY (user_id) REFERENCES users (id)
                        )
                    ''')
                    
                    # 复制数据（不复制type字段）
                    self.cursor.execute('''
                        INSERT INTO information (id, user_id, title, content, category, create_time, update_time)
                        SELECT id, user_id, title, content, category, create_time, update_time
                        FROM information_old
                    ''')
                    
                    # 删除旧表
                    self.cursor.execute("DROP TABLE information_old")
                    self.conn.commit()
                except Exception as inner_e:
                    # 如果重命名表失败，尝试其他方法
                    try:
                        # 尝试添加默认值以移除NOT NULL约束
                        self.cursor.execute("UPDATE information SET type = '' WHERE type IS NULL")
                        self.conn.commit()
                    except:
                        pass
        except Exception as e:
            # 记录错误但不阻止程序继续运行
            pass
    
    def close(self):
        """关闭数据库连接"""
        if self.conn:
            self.conn.close()
    
    def commit(self):
        """提交事务"""
        if self.conn:
            self.conn.commit()
    
    def rollback(self):
        """回滚事务"""
        if self.conn:
            self.conn.rollback()